import os
import shutil
import tempfile
from util_slash import f_win, f_unix

def get_total_size(src_dir):
    """
    计算目录的总大小（包括所有子目录和文件）

    参数:
        src_dir (str): 源目录路径

    返回:
        int: 总字节数
    """
    total_size = 0
    if os.path.isfile(src_dir):
        total_size += os.path.getsize(src_dir)
        return total_size

    for dirpath, dirnames, filenames in os.walk(src_dir):
        for filename in filenames:
            filepath = os.path.join(dirpath, filename)
            # 跳过符号链接以免重复计算
            if not os.path.islink(filepath):
                try:
                    total_size += os.path.getsize(filepath)
                except (OSError, FileNotFoundError):
                    # 处理无法访问的文件
                    pass
    return total_size


def replace_directory_with_copy(source_dir, target_dir, chunk_size=1024 * 1024, fake_copy=False, real_delete=True):
    """
    用源目录完全替换目标目录：
    1. 删除目标目录（如果存在）及其所有内容
    2. 将源目录及其所有内容复制到目标目录

    参数:
        source_dir (str): 源目录路径
        target_dir (str): 目标目录路径
        chunk_size (int): 复制文件时使用的块大小（字节），默认1MB
    """
    if fake_copy:
        print(f"\033[46m[Manual Copy]\033[0m from: {f_win(source_dir)} to: {f_win(target_dir)}")
        return True

    # 检查target_dir是否安全可以执行删除操作
    if target_dir.replace("\\", "/").split('/')[-1] != source_dir.replace("\\", "/").split("/")[-1]:
        raise AssertionError(f"目标目录不安全，请检查配置: {target_dir}")

    # 验证源目录是否存在
    if not os.path.exists(source_dir):
        print(f"源目录 '{source_dir}' 不存在")
        raise FileNotFoundError(f"源目录 '{source_dir}' 不存在")

    # 验证源目录确实是一个目录
    # if not os.path.isdir(source_dir):
    #     print(f"'{source_dir}' 不是一个目录")
    #     raise NotADirectoryError(f"'{source_dir}' 不是一个目录")

    try:
        # 第一步：删除目标目录（如果存在）
        if os.path.exists(target_dir):
            tmp_path = tempfile.mktemp(dir=os.path.dirname(target_dir))
            print(f"目标目录 '{target_dir}' 重命名为 '{tmp_path}'")
            shutil.move(target_dir, tmp_path)  # shutil.move 无论对于目录还是文件都可以用
            if real_delete:
                if os.path.isdir(source_dir):
                    print(f"正在删除目标目录 '{tmp_path}' 及其所有内容...")
                    shutil.rmtree(tmp_path)  # 递归删除目录及其所有内容[1,2](@ref)
                    print(f"目标目录 '{tmp_path}' 已成功删除")
                else:
                    print(f"正在删除文件 '{tmp_path}' ...")
                    os.remove(tmp_path)
                    print(f"目标文件 '{tmp_path}' 已成功删除")
        else:
            print(f"目标目录 '{target_dir}' 不存在，无需删除")

        # 第二步：计算源目录总大小（用于进度条）
        print("正在计算源目录总大小...")
        total_size = get_total_size(source_dir)

        if total_size == 0:
            print("源目录中没有文件可复制或所有文件大小为0。")
            # 即使没有文件，也创建空目录结构
            shutil.copytree(source_dir, target_dir, dirs_exist_ok=True)
            return True

        print(f"源目录总大小: {total_size} 字节 ({total_size / (1024 * 1024):.2f} MB)")

        if os.path.isdir(source_dir):  # 复制目录
            # 第三步：创建目标目录的根目录
            os.makedirs(target_dir, exist_ok=True)
            from tqdm import tqdm
            # 第四步：使用tqdm创建进度条，基于数据量进行复制
            with tqdm(
                    total=total_size,
                    unit='B',
                    unit_scale=True,
                    unit_divisor=1024,
                    desc=f"复制 {os.path.basename(source_dir)}",
                    ncols=80
            ) as pbar:
                for dirpath, dirnames, filenames in os.walk(source_dir):
                    # 计算当前目录相对于源目录的相对路径
                    rel_path = os.path.relpath(dirpath, source_dir)
                    # 构建目标目录路径
                    dest_dirpath = os.path.join(target_dir, rel_path)

                    # 确保目标子目录存在
                    if not os.path.exists(dest_dirpath):
                        os.makedirs(dest_dirpath, exist_ok=True)

                    # 复制当前目录中的所有文件
                    for filename in filenames:
                        src_file_path = os.path.join(dirpath, filename)
                        dest_file_path = os.path.join(dest_dirpath, filename)

                        try:
                            # 使用分块复制以便更新进度条
                            with open(src_file_path, 'rb') as src_file:
                                with open(dest_file_path, 'wb') as dest_file:
                                    while True:
                                        chunk = src_file.read(chunk_size)
                                        if not chunk:
                                            break
                                        dest_file.write(chunk)
                                        # 更新进度条
                                        pbar.update(len(chunk))
                                        # 更新进度条后缀信息，显示当前正在复制的文件
                                        pbar.set_postfix(file=filename[:20] + '...' if len(filename) > 20 else filename)

                            # 保留文件元数据（修改时间、权限等）
                            shutil.copystat(src_file_path, dest_file_path)

                        except Exception as copy_error:
                            print(f"复制文件 '{src_file_path}' 到 '{dest_file_path}' 失败: {copy_error}")
                            # 可以选择继续复制其他文件或抛出异常

            print(f"成功将源目录 '{source_dir}' 复制到 '{target_dir}'")
            return True
        else:  # 复制文件
            from tqdm import tqdm
            with tqdm(
                    total=total_size,
                    unit='B',
                    unit_scale=True,
                    unit_divisor=1024,
                    desc=f"复制 {os.path.basename(source_dir)}",
                    ncols=80
            ) as pbar:
                filename = os.path.basename(source_dir)
                try:
                    # 使用分块复制以便更新进度条
                    with open(source_dir, 'rb') as src_file:
                        with open(target_dir, 'wb') as dest_file:
                            while True:
                                chunk = src_file.read(chunk_size)
                                if not chunk:
                                    break
                                dest_file.write(chunk)
                                # 更新进度条
                                pbar.update(len(chunk))
                                # 更新进度条后缀信息，显示当前正在复制的文件
                                pbar.set_postfix(file=filename[:20] + '...' if len(filename) > 20 else filename)

                except Exception as copy_error:
                    print(f"复制文件 '{source_dir}' 到 '{target_dir}' 失败: {copy_error}")
                    # 可以选择继续复制其他文件或抛出异常
            return True
    except PermissionError as e:
        print(f"权限错误: {e}")
        raise
    except FileExistsError as e:
        print(f"文件已存在错误: {e}")
        raise
    except Exception as e:
        print(f"操作过程中发生未知错误: {e}")
        raise


if __name__ == "__main__":
    # source_dir = r"\\diskstation3\10T_01\Fitgirl\Cyberpunk 2077 [FitGirl Repack]"
    # target_dir = r"\\diskstation3\14T_01\fitgirl_10T_01\Cyberpunk 2077 [FitGirl Repack]"
    source_dir = r"E:/civ5bnw_103276_cn_20141028.rar"
    target_dir = r"F:/civ5bnw_103276_cn_20141028.rar"
    replace_directory_with_copy(source_dir, target_dir)
